import sapien.core as sapien
from sapien.utils.viewer import Viewer

import mplib
from mplib import Pose


class MPLibPlanner:

    def __init__(self, scene, urdf, srdf=None, move_group='tool0'):
        self._scene = scene
        self.planner = mplib.Planner(
            urdf=urdf,
            srdf=srdf,
            move_group=move_group,
        )

    def move_to_pose_with_RRTConnect(self, pose: Pose):
        """
        Plan and follow a path to a pose using RRTConnect

        Args:
            pose: mplib.Pose
        """
        # result is a dictionary with keys 'status', 'time', 'position', 'velocity',
        # 'acceleration', 'duration'
        # plan_pose ankor
        print("plan_pose")
        result = self.planner.plan_pose(pose, self.robot.get_qpos(), time_step=1 / 250)
        # plan_pose ankor end
        if result["status"] != "Success":
            print(result["status"])
            return -1
        return 0

    def move_to_pose_with_screw(self, pose):
        """
        Interpolative planning with screw motion.
        Will not avoid collision and will fail if the path contains collision.
        """
        result = self.planner.plan_screw(
            pose,
            self.robot.get_qpos(),
            time_step=1 / 250,
        )
        if result["status"] == "Success":
            return 0
        else:
            # fall back to RRTConnect if the screw motion fails (say contains collision)
            return self.move_to_pose_with_RRTConnect(pose)

    def move_to_pose(self, pose, with_screw=True):
        """API to multiplex between the two planning methods"""
        if with_screw:
            return self.move_to_pose_with_screw(pose)
        else:
            return self.move_to_pose_with_RRTConnect(pose)
        
    def forward_kinematic(self, qpos):
        self.planner.pinocchio_model.compute_forward_kinematics(qpos)
        move_link_idx = self.planner.link_name_2_idx[self.planner.move_group]
        return self.planner.pinocchio_model.get_link_pose(move_link_idx)
    
    def forward_velocity(self, qpos, qvel):
        self.planner.pinocchio_model.compute_forward_kinematics(qpos)
        self.planner.pinocchio_model.compute_full_jacobian(qpos)
        move_link_idx = self.planner.link_name_2_idx[self.planner.move_group]
        J = self.planner.pinocchio_model.get_link_jacobian(move_link_idx, local=False)

        return J @ qvel